{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc1b5843-eca4-47f5-8e3c-4e4dc83775f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "from functools import lru_cache\n",
    "import pandas as pd\n",
    "from gs_quant.api.gs.assets import GsAssetApi\n",
    "from gs_quant.data import Dataset\n",
    "from gs_quant.session import GsSession, Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde1622d-5e57-4f4e-9602-7ac6aa75f23a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Internals\n",
    "@lru_cache(maxsize=128)\n",
    "def _get_coverage(ds):\n",
    "    print(\"Fetching coverage\", end=\" \")\n",
    "    cov = ds.get_coverage()[\"assetId\"].to_list()\n",
    "    print(\"[DONE]\")\n",
    "    return cov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34341f3d-cf64-4519-b808-335d874b3cd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _nest_list(flat_list, nested_size):\n",
    "    return [flat_list[i:i + nested_size] for i in range(0, len(flat_list), nested_size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f73538c6-a774-4ae4-bacd-fa584428cd9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_data_by_asset_parameters(ds, start_date, end_date, pricing_location=\"NYC\", chunk_size=100, **kwargs):\n",
    "    coverage = _get_coverage(ds)\n",
    "    assets = []\n",
    "    asset_chunks = _nest_list(coverage, 5000)\n",
    "    print(\"Reading Assets:\", end=\" \")\n",
    "    for ch in asset_chunks:\n",
    "        print(\"#\", end=\"\")\n",
    "        kwargs[\"id\"] = ch\n",
    "        assets += GsAssetApi.get_many_assets([\"id\"], limit=10000, **kwargs)\n",
    "    print(\" [DONE]\")\n",
    "\n",
    "    ids = [a.id for a in assets]\n",
    "    print(\"Assets to fetch \" + str(len(ids)))\n",
    "    chunks = _nest_list(ids, chunk_size)\n",
    "    frames = [pd.DataFrame()]\n",
    "    print(\"Reading Data:\", end=\" \")\n",
    "    for ch in chunks:\n",
    "        print(\"#\", end=\"\")\n",
    "        frames.append(ds.get_data(start_date, end_date, assetId=ch, pricingLocation=pricing_location))\n",
    "    print(\" [DONE]\")\n",
    "    return pd.concat(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93546060-b2a4-42c1-94f9-87939b07e1e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_fxivol(ds, bbid: str = None, tenor=None, strike_reference=None, relative_strike=None, pricing_location=\"NYC\",\n",
    "               start_date=datetime.date.today(), end_date=datetime.date.today()):\n",
    "    \"\"\"\n",
    "    :param ds: Dataset object\n",
    "    :param bbid: selected identifier i.e. EURUSD\n",
    "    :param tenor: tenor of the option i.e. 1y, 5y\n",
    "    :param strike_reference: one of \"delta\", \"forward\", \"spot\", one supported at at time\n",
    "    :param relative_strike: moneyness of the option i.e. -5, 0, 5 ... etc Mixing negative and positive are not supported\n",
    "    :param pricing_location:\n",
    "    :param start_date: start date of the query\n",
    "    :param end_date:  end date of the query\n",
    "    :return: DataFrame with data\n",
    "    \"\"\"\n",
    "    query = dict()\n",
    "    if bbid is not None:\n",
    "        query[\"asset_parameters_call_currency\"] = [bbid[0:3], bbid[3:]]\n",
    "        query[\"asset_parameters_put_currency\"] = [bbid[0:3], bbid[3:]]\n",
    "    if strike_reference is not None:\n",
    "        relative_strike_l = relative_strike\n",
    "        if strike_reference == \"delta\" and relative_strike is not None:\n",
    "            if not isinstance(relative_strike, list):\n",
    "                relative_strike_l = [relative_strike]\n",
    "            if min(relative_strike_l) < 0 < max(relative_strike_l):\n",
    "                print(\"ERROR: Relative strikes should be same sign in the query to work properly\")\n",
    "                return None\n",
    "            query[\"asset_parameters_option_type\"] = \"Put\" if relative_strike_l[0] < 0 else \"Call\"\n",
    "            query[\"asset_parameters_strike_price_relative\"] = [str(abs(s)) + \"D\" if s != 0 else \"DN\" for s in\n",
    "                                                               relative_strike_l]\n",
    "\n",
    "    if strike_reference == \"forward\":\n",
    "        query[\"asset_parameters_strike_price_relative\"] = \"ATMF\"\n",
    "    if strike_reference == \"spot\":\n",
    "        query[\"asset_parameters_strike_price_relative\"] = [\"ATMS\", \"Spot\", \"ATM\"]\n",
    "\n",
    "    if tenor is not None:\n",
    "        query[\"asset_parameters_expiration_date\"] = tenor\n",
    "\n",
    "    print(query)\n",
    "    return _get_data_by_asset_parameters(ds, start_date, end_date, pricing_location=pricing_location, **query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "516cd795-dcf1-40ee-bf59-cb7189d3a204",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_fx_fwd(ds, bbid: str = None, tenor=None, pricing_location=\"NYC\",\n",
    "               start_date=datetime.date.today(), end_date=datetime.date.today()):\n",
    "    \"\"\"\n",
    "    :param ds: Dataset object\n",
    "    :param bbid: selected identifier i.e. EURUSD\n",
    "    :param tenor: tenor of the option i.e. 1y, 5y\n",
    "    :param pricing_location:\n",
    "    :param start_date: start date of the query\n",
    "    :param end_date:  end date of the query\n",
    "    :return: DataFrame with data\n",
    "    \"\"\"\n",
    "    args = dict()\n",
    "    if bbid is not None:\n",
    "        args[\"asset_parameters_pair\"] = bbid\n",
    "    if tenor is not None:\n",
    "        args[\"asset_parameters_expiration_date\"] = tenor\n",
    "    return _get_data_by_asset_parameters(ds, start_date, end_date, pricing_location=pricing_location, **args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28f5b009-3302-4829-b2a3-d6f56c61577e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# USER DETAILS\n",
    "CLIENT_ID = None\n",
    "SECRET = None\n",
    "GsSession.use(Environment.PROD, client_id=CLIENT_ID, client_secret=SECRET)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5739dad6-aee6-4dde-b504-a355233fc4cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#SELECT DATASET\n",
    "FX_VOL_DS = Dataset(\"FXIVOL_V2_PREMIUM\")\n",
    "start_date = datetime.date(2021, 12, 17)\n",
    "end_date = datetime.date(2022, 12, 17)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3528c97b-9f4a-4381-a92f-5bf71a3d304a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = get_fxivol(FX_VOL_DS, start_date=start_date, end_date=end_date, bbid=\"EURUSD\", relative_strike=[-10],tenor=\"5y\",\n",
    "                     strike_reference=\"delta\")\n",
    "print(vol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "992630f6-7614-4165-94f8-18ebca34e25e",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = get_fxivol(FX_VOL_DS, start_date=start_date, end_date=end_date, bbid=\"EURUSD\", relative_strike=[10],tenor=\"5y\",\n",
    "                     strike_reference=\"delta\")\n",
    "print(vol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5730f24-8624-44c8-93a5-5cf4683b2423",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = get_fxivol(FX_VOL_DS, start_date=start_date, end_date=end_date, bbid=\"USDPLN\",\n",
    "                     strike_reference=\"delta\",\n",
    "                     tenor=\"3m\")\n",
    "print(vol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d62c755-55c9-42e7-9407-9f78ab1f10b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = get_fxivol(FX_VOL_DS, start_date=start_date, end_date=end_date, bbid=\"EURUSD\", relative_strike=[0, 10, 15],\n",
    "                     strike_reference=\"delta\",\n",
    "                     tenor=\"1y\")\n",
    "print(vol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac617f52-0004-487a-b6a2-e67757846a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = get_fxivol(FX_VOL_DS, start_date=start_date, end_date=end_date, bbid=\"USDPLN\",\n",
    "                     strike_reference=\"forward\",\n",
    "                     tenor=\"5y\")\n",
    "print(vol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4637647-6cd0-4695-b081-8a22b75081bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol = get_fxivol(FX_VOL_DS, start_date=start_date, end_date=end_date, bbid=\"USDPLN\",\n",
    "                     strike_reference=\"spot\",\n",
    "                     tenor=\"1y\")\n",
    "print(vol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba12394-7b8d-4136-974e-5d100ad4c091",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
